import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

def plot_figure3_from_data_final():
    """
    Reproduces Figure 3 from the paper with high fidelity, matching the
    color scheme (black base + red layers) and visual style of the original.
    """
    # 1. Store the provided data in a pandas DataFrame
    data = {
        'λ': ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5', '0.6', '0.7', '0.8', '0.9', '1.0'],
        '0': [92.3, 93, 93.4, 94.3, 95, 95.6, 95.9, 96, 96.1, 96.1, 96.3],
        '0.05': [90.7, 91.9, 92.7, 93, 94.4, 94.8, 94.7, 94.6, 94.3, 92.9, 90.3],
        '0.1': [90.8, 91.4, 92.4, 93.1, 93.4, 94.1, 93.9, 93.4, 91.5, 87.7, 77.8],
        '0.15': [90.7, 91.3, 91.7, 92.3, 92.9, 93.1, 93.9, 91.7, 88.1, 80.3, 62.2],
        '0.2': [89.6, 90.1, 91.1, 91.8, 92, 92.5, 93.3, 88.8, 84.1, 72.1, 41.3],
        '0.25': [88.2, 89.2, 90.1, 90.7, 90.1, 90.6, 88.8, 85.8, 79.5, 59.4, 20.9],
        '0.3': [86.6, 87.7, 87.7, 88.2, 88.4, 87.9, 85.3, 82.2, 72.9, 47.5, 8.4]
    }
    df = pd.DataFrame(data)
    df = df.set_index('λ')
    
    # Add the 'fine-tuning' data, estimated visually from the paper's Figure 3
    fine_tuning_data = pd.Series({
        '0': 98.0, '0.05': 95.0, '0.1': 90.0, '0.15': 85.0, 
        '0.2': 80.0, '0.25': 72.0, '0.3': 62.0
    }, name='fine-tuning')
    df.loc['fine-tuning'] = fine_tuning_data

    budgets = ['0', '0.05', '0.1', '0.15', '0.2', '0.25', '0.3']
    
    # 2. Calculate the differential values for the stacked segments
    plot_data = pd.DataFrame(index=df.index)
    plot_data['0.3'] = df['0.3'] # The base of the stack is the accuracy at the highest budget
    for i in range(len(budgets) - 2, -1, -1):
        current_budget_col = budgets[i+1]
        higher_budget_col = budgets[i]
        # The height of each subsequent layer is the difference in accuracy
        plot_data[higher_budget_col] = df[higher_budget_col] - df[current_budget_col]

    # 3. Create the plot with the corrected color scheme
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Final color map, using black for the base as in the original paper
    color_map = {
        '0.3': '#000000',      # Black base (most robust)
        '0.25': '#500000',     # Darkest Red
        '0.2': '#900000',      
        '0.15': '#D00000',     # Bright Red
        '0.1': '#FF5050',
        '0.05': '#FF9090',     # Light Red
        '0': '#FFD0D0'       # Lightest Pink
    }
    
    bottom = np.zeros(len(df))
    bar_width = 0.8
    
    # Plot bars layer by layer, from the black base up to the lightest top layer
    for budget in reversed(budgets):
        values = plot_data[budget].values
        ax.bar(df.index, values, bottom=bottom, width=bar_width, label=budget, color=color_map[budget])
        bottom += values
        
    # 4. Customize labels, title, and ticks
    ax.set_xlabel('λ', fontsize=16)
    ax.set_ylabel('Accuracy (%)', fontsize=16)
    ax.tick_params(axis='both', which='major', labelsize=14)
    ax.set_ylim(0, 101)
    
    # Remove title and top/right spines for a cleaner look, closer to the paper
    ax.set_title('') 
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    
    # 5. Create a custom legend, moved slightly outside the plot area
    legend_patches = [mpatches.Patch(color=color_map[b], label=b) for b in budgets]
    ax.legend(
        handles=legend_patches, 
        title='Budget', 
        loc='upper left', 
        bbox_to_anchor=(1.01, 1.0), 
        fontsize=12, 
        title_fontsize=13, 
        frameon=True, 
        edgecolor='black'
    )
    
    # Use tight_layout and then adjust subplot parameters for a perfect fit
    plt.tight_layout()
    plt.subplots_adjust(right=0.85) # Make space for the legend on the right
    
    # Save the final figure
    output_filename = 'figure3_reproduction_final.png'
    plt.savefig(output_filename, dpi=300)
    print(f"Final plot saved as {output_filename}")
    
    plt.show()

if __name__ == '__main__':
    plot_figure3_from_data_final() 